`include "defines.v"

module IF_Stage (
    input  wire                       clock,
    input  wire                       reset,

    input  wire                       stall,
    input  wire                       flush,

    input  wire                       exception_jump_flag,
    input  wire [`BUSLEN-1:0]         exception_jump_addr,

    input  wire                       inst_jump_or_branch_taken_flag,
    input  wire [`BUSLEN-1:0]         inst_jump_or_branch_taken_addr,

    output wire                       pc_valid,
    output wire                       inst_fetched,
    output wire [`INSTLEN-1:0]        inst,
    output reg  [`BUSLEN-1:0]         pc,

    // To axi_rw
    output wire                       if_r_req,
    input  wire                       if_r_okay,
    input  wire                       if_r_handshaked,
    input  wire [`AXI_DATA_WIDTH-1:0] if_data_read,
    output wire [`AXI_ADDR_WIDTH-1:0] if_r_addr,
    output wire [ 1: 0]               if_r_size,
    input  wire [ 1: 0]               if_r_resp,

    output wire                       inst_valid
);


    //----------Instruction Fetching----------//
    reg  prefetching_enable;
    reg  prefetching_enable_en;
    always @(posedge clock) begin
        if(reset) begin
            prefetching_enable <= 1'b0;
            prefetching_enable_en <= 1'b1;
        end
        else if(if_r_req & if_r_handshaked & prefetching_enable_en) begin
            prefetching_enable <= pc[31];
            prefetching_enable_en <= 1'b0;
        end
        else if(if_r_okay) begin
            prefetching_enable_en <= 1'b1;
        end
    end

    // Instruction Prefetch
    reg  inst_prefetched;
    wire inst_prefetched_en;
    wire inst_prefetched_rst;

    assign inst_prefetched_en  = (inst_fetched | inst_fetched_keep) & ~stall & prefetching_enable;
    assign inst_prefetched_rst = inst_prefetched & ~stall | flush;

    always @(posedge clock) begin
        if(reset) inst_prefetched <= 0;
        else if(inst_prefetched_rst) inst_prefetched <= 0;
        else if(inst_prefetched_en)  inst_prefetched <= 1'b1;
    end

    // Instruction Fetching
    assign if_r_req  = ~inst_fetched & ~inst_fetched_keep & ~inst_prefetched;
    assign if_r_addr = pc;
    assign if_r_size = `SIZE_D;

    reg  inst_fetched_keep;
    wire inst_fetched_keep_en;
    wire inst_fetched_keep_rst;

    assign inst_fetched_keep_en  = inst_fetched & stall & prefetching_enable;
    assign inst_fetched_keep_rst = inst_fetched_keep & ~stall | flush;
    
    always @(posedge clock) begin
        if(reset) inst_fetched_keep <= 0;
        else if(inst_fetched_keep_rst) inst_fetched_keep <= 0;
        else if(inst_fetched_keep_en)  inst_fetched_keep <= 1'b1;
    end

    // Instruction Fetched Discard Mechanism
    reg discard;

    always@(posedge clock) begin
        if(reset)          discard <= 1'b0;
        else if(if_r_okay) discard <= 1'b0;
        else if(flush)     discard <= if_r_handshaked;
        else               discard <= discard;
    end

    assign inst_fetched = if_r_okay & ~discard & (if_r_resp == 0);
    assign inst =  (inst_fetched | inst_fetched_keep)  ?  if_data_read[31:0]  :  
                   inst_prefetched                     ?  if_data_read[63:32]  :  32'h13  ;
    assign pc_valid = 1'b1;
    assign inst_valid = (inst_fetched | inst_fetched_keep | inst_prefetched) & ~stall;

    //----------PC Generation----------//
    wire [2:0] pc_src; /* 000:+4; 001:hold; 010:jump or branch taken; 100:exception; */

    assign pc_src[2] = exception_jump_flag;
    assign pc_src[1] = inst_jump_or_branch_taken_flag;
    assign pc_src[0] = stall | ~(inst_fetched | inst_fetched_keep | inst_prefetched);

    always@(posedge clock)
    begin
        if(reset) pc <=`PC_START;
        else      pc <= pc_src[2] ?  exception_jump_addr             :
                        pc_src[1] ?  inst_jump_or_branch_taken_addr  :
                        pc_src[0] ?  pc                              :  pc+64'h4 ;
    end

endmodule